#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Float64MultiArray

# --- 定义 Leader 和 Follower 的关节名称 ---
# 必须与 'so101_controllers.yaml' 和 'unified_joy_controller.py' 中的定义一致
ARM_JOINT_NAMES = ["1", "2", "3", "4", "5"]
GRIPPER_JOINT_NAME = "6"

# --- 定义 Leader 和 Follower 的 ROS 主题 ---
# !!! 重要 !!!
# 您需要将 LEADER_JOINT_STATES_TOPIC 替换为您的 Leader 机械臂
# 发布其关节状态的主题名称。
#
# 假设 Leader 机械臂也使用 /joint_states，但为了避免冲突，
# 更好的做法是在启动 Leader 的 ros2_control 时，为其指定一个
# 独立的 namespace，例如 "leader"。
# 这样，它的主题就会是 "/leader/joint_states"。
#
# 如果 Leader 和 Follower 在同一个 /joint_states 主题上发布，
# 这个节点将会收到混合的数据，逻辑会变得非常复杂。
#
# *** 请在此处确认并修改 Leader 的主题名称 ***
LEADER_JOINT_STATES_TOPIC = "/leader/joint_states" # <--- 示例值, 请修改!

# Follower 的指令主题 (根据您的 unified_joy_controller.py 确认)
FOLLOWER_ARM_COMMAND_TOPIC = "/arm_controller/commands"
FOLLOWER_GRIPPER_COMMAND_TOPIC = "/gripper_controller/commands"


class LeaderFollowerBridge(Node):
    def __init__(self):
        super().__init__('leader_follower_bridge')

        # --- Follower 的发布者 ---
        self.arm_cmd_pub = self.create_publisher(
            Float64MultiArray, FOLLOWER_ARM_COMMAND_TOPIC, 10
        )
        self.gripper_cmd_pub = self.create_publisher(
            Float64MultiArray, FOLLOWER_GRIPPER_COMMAND_TOPIC, 10
        )

        # --- Leader 的订阅者 ---
        self.leader_joint_sub = self.create_subscription(
            JointState,
            LEADER_JOINT_STATES_TOPIC,
            self.leader_state_callback,
            10
        )

        # 用于存储 Leader 关节名称 -> 索引 的映射
        self.leader_joint_map = {}
        self.arm_joint_indices = []
        self.gripper_joint_index = -1

        self.get_logger().info("✅ Leader-Follower 桥接节点已启动。")
        self.get_logger().info(f"    监听 Leader 状态于: {LEADER_JOINT_STATES_TOPIC}")
        self.get_logger().info(f"    转发手臂指令至: {FOLLOWER_ARM_COMMAND_TOPIC}")
        self.get_logger().info(f"    转发夹爪指令至: {FOLLOWER_GRIPPER_COMMAND_TOPIC}")
        self.get_logger().warn(f"*** 请确保 Leader 机械臂正在 '{LEADER_JOINT_STATES_TOPIC}' 上发布状态! ***")


    def leader_state_callback(self, msg: JointState):
        """
        接收 Leader 的状态, 提取位置, 并作为指令发布给 Follower。
        """

        # --- 步骤 1: (仅首次) 构建关节名称到索引的映射 ---
        if not self.leader_joint_map:
            self.get_logger().info("首次接收到 Leader 状态, 正在构建关节索引映射...")
            self.leader_joint_map = {name: i for i, name in enumerate(msg.name)}
            
            try:
                # 查找手臂关节的索引
                self.arm_joint_indices = [self.leader_joint_map[name] for name in ARM_JOINT_NAMES]
                
                # 查找夹爪关节的索引
                self.gripper_joint_index = self.leader_joint_map[GRIPPER_JOINT_NAME]
                
                self.get_logger().info("✅ 关节索引映射构建成功。")
                self.get_logger().info(f"    手臂关节 (1-5) 索引: {self.arm_joint_indices}")
                self.get_logger().info(f"    夹爪关节 (6) 索引: {self.gripper_joint_index}")

            except KeyError as e:
                self.get_logger().error(f"!!! 严重错误: 在 Leader 的 joint_states 中未找到关节 {e}。")
                self.get_logger().error("    请检查 LEADER_JOINT_STATES_TOPIC 和 ARM/GRIPPER_JOINT_NAMES 是否正确。")
                self.get_logger().error(f"    Leader 报告的关节有: {list(msg.name)}")
                # 禁用回调以防止进一步错误
                self.leader_joint_sub.destroy() 
                return
        
        # --- 步骤 2: 提取手臂位置并发布 ---
        if self.arm_joint_indices:
            try:
                arm_positions = [msg.position[i] for i in self.arm_joint_indices]
                arm_cmd_msg = Float64MultiArray()
                arm_cmd_msg.data = arm_positions
                self.arm_cmd_pub.publish(arm_cmd_msg)
            except IndexError:
                self.get_logger().warn("Leader joint_states 消息长度与预期不符 (手臂)。")


        # --- 步骤 3: 提取夹爪位置并发布 ---
        if self.gripper_joint_index != -1:
            try:
                gripper_position = msg.position[self.gripper_joint_index]
                gripper_cmd_msg = Float64MultiArray()
                gripper_cmd_msg.data = [gripper_position] # 夹爪控制器只需要一个值的数组
                self.gripper_cmd_pub.publish(gripper_cmd_msg)
            except IndexError:
                 self.get_logger().warn("Leader joint_states 消息长度与预期不符 (夹爪)。")


def main(args=None):
    rclpy.init(args=args)
    bridge_node = LeaderFollowerBridge()
    rclpy.spin(bridge_node)
    bridge_node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()